# 动态规划50题 https://www.bilibili.com/video/BV1aa411f7uT
# 13/50 寻找数组中心下标
# leetcode第724题: https://leetcode.cn/problems/find-pivot-index/description/
# Date: 2024/11/1

from leetcode import test


def pivotIndex(nums: list[int]) -> int:
    """我的方法，使用动态规划来求前i项和，如果相邻两项a、b之和等于最后一项（所有元素之和），则中心下标即为a"""
    n = len(nums)
    if n < 2:
        return 0

    dp = [0] * (n + 1)
    dp[0] = 0
    for i in range(1, n + 1):
        dp[i] = dp[i - 1] + nums[i - 1]

    for i in range(n):
        if dp[n] - dp[i] - dp[i + 1] == 0:
            return i
    return -1


if __name__ == '__main__':
    inp = [{"nums": [1, 7, 3, 6, 5, 6]}, {"nums": [1, 2, 3]}, {"nums": [2, 1, -1]}]
    out = [3, -1, 0]
    test.test_function(pivotIndex, inp, out)
